import numpy as np
import gurobipy as gp
from gurobipy import GRB

import verify_piecewise_utils
from NN import cdinn2
from prob import *

# case_name = "zhrC15"

margin1 = 0.01
margin2 = 0.01

# https://support.gurobi.com/hc/en-us/articles/4414392016529-How-do-I-model-conditional-statements-in-Gurobi
def relu_der(m, x, useVars):
    eps = 0.0001
    M = 10000 + eps

    if useVars:
        relu_der_x = m.addMVar(x.shape, vtype=GRB.CONTINUOUS)
        for i in range(x.shape[0]):
            b = m.addVar(vtype=GRB.BINARY)
            m.addConstr(x[i] >= eps - M * (1 - b))
            m.addConstr(x[i] <= M * b)
            m.addConstr((b == 1) >> (relu_der_x[i] == 1))
            m.addConstr((b == 0) >> (relu_der_x[i] == 0))
        m.update()
    else:
        relu_der_x = []
        for i in range(x.shape[0]):
            if x[i] > 0:
                relu_der_x.append(1)
            else:
                relu_der_x.append(0)
        relu_der_x = np.array(relu_der_x)
    return relu_der_x


def barrier_f(m, x0, wb_list, wb_length, useVars):
    x_layer = x0
    if useVars:
        jacobian = np.eye(x0.shape[0])
    else:
        jacobian = np.eye(len(x0))
    for i in range(wb_length):
        W = wb_list[i * 2]
        b = wb_list[i * 2 + 1]
        if useVars:
            y_layer = m.addMVar((W.shape[0],), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)
            E = np.identity(y_layer.shape[0])
            expr = W @ x_layer + b - E @ y_layer
            m.addConstr(expr == 0)
        else:
            # print("W", W.shape)
            # x_layer = np.array(x_layer)
            # print("x_layer", x_layer.shape)
            # print("b:", b.shape)
            y_layer = W @ x_layer + b
            # y_layer = np.array(y_layer)
            # print("y_layer:", y_layer.shape)
        jacobian = W @ jacobian

        if i != wb_length - 1:
            if useVars:
                z_layer = m.addMVar((y_layer.shape[0],), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)
                # z = activate(y)
                for j in range(y_layer.shape[0]):
                    m.addConstr(z_layer[j] == gp.max_(y_layer[j], 0.0))
            else:
                z_layer = y_layer
                for j in range(y_layer.shape[0]):
                    z_layer[j] = max(y_layer[j], 0.0)
            x_layer = z_layer

            array_relu_der = []
            relu_der_f = relu_der(m, y_layer, useVars)
            for j in range(relu_der_f.shape[0]):
                array_relu_der.append(relu_der_f[j])
            jacobian = np.diag(array_relu_der) @ jacobian
        else:
            x_layer = y_layer

        if useVars:
            m.update()

    return x_layer, jacobian[0]


def verify_initCond(wb_1_list, wb_2_list, wb_length):
    cex_init = []
    res_init = []
    m = gp.Model()
    m.setParam('Outputflag', 0)

    x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)
    for i in range(DIM):
        m.addConstr(x0[i] >= init_min[i])
        m.addConstr(x0[i] <= init_max[i])

    # if init_shape == 2:
    #     p0 = (init_min[0] + init_max[0]) / 2.0
    #     p1 = (init_min[1] + init_max[1]) / 2.0
    #     p2 = (init_min[2] + init_max[2]) / 2.0
    #     r = p0 - init_min[0]
    #     m.addConstr((x0[0] - p0) * (x0[0] - p0) + (x0[1] - p1) * (x0[1] - p1) + (x0[2] - p2) * (x0[2] - p2) <= r * r)

    # f1_x0, df1dx_x0 = barrier_f(m, x0, wb_1_list, wb_length, True)
    f2_x0, df2dx_x0 = barrier_f(m, x0, wb_2_list, wb_length, True)

    # xt = np.random.randn(DIM)
    xt = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    result_pre = 0

    while True:
        # print("================")
        f1_xt, df1dx_xt = barrier_f(m, xt, wb_1_list, wb_length, False)
        # f2_xt, df2dx_xt = barrier_f(m, xt, wb_2_list, wb_length, False)

        # print("f1_xt:", f1_xt)
        # print("df1dx_xt[0]:", df1dx_xt)
        # print("f2_x0:", f2_x0)

        m.setObjective(f1_xt
                       + df1dx_xt[0] * (x0[0] - xt[0]) + df1dx_xt[1] * (x0[1] - xt[1]) + df1dx_xt[2] * (x0[2] - xt[2])
                       + df1dx_xt[3] * (x0[3] - xt[3]) + df1dx_xt[4] * (x0[4] - xt[4]) + df1dx_xt[5] * (x0[5] - xt[5])
                       + df1dx_xt[6] * (x0[6] - xt[6]) + df1dx_xt[7] * (x0[7] - xt[7]) + df1dx_xt[8] * (x0[8] - xt[8])
                       + df1dx_xt[9] * (x0[9] - xt[9])
                       - f2_x0,
                       GRB.MAXIMIZE)
        m.optimize()
        # m.computeIIS()
        # m.write("model1.ilp")
        # print("m.status:", m.status)

        # print("x0.x:", x0.x)
        # print("m.objVal:", m.objVal)
        if abs(m.objVal - result_pre) < margin1 or np.linalg.norm(x0.x - xt) < margin2:
            break
        xt = x0.x
        # print("xt[0]:", xt[0])
        # print("xt[1]:", xt[1])
        result_pre = m.objVal

    # print("maximize init value:", m.objVal)  # -0.07222615058680137
    if m.objVal > 0:
        cex_init.append(np.array(x0.x))
        res_init.append(m.objVal)
    return cex_init, res_init


def verify_unsafeCond(wb_1_list, wb_2_list, wb_length):
    cex_unsafe = []
    res_unsafe = []
    m = gp.Model()
    m.setParam('Outputflag', 0)

    x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf)
    for i in range(DIM):
        m.addConstr(x0[i] >= unsafe_min[i])
        m.addConstr(x0[i] <= unsafe_max[i])

    # if unsafe_shape == 2:
    #     p0 = (unsafe_min[0] + unsafe_max[0]) / 2.0
    #     p1 = (unsafe_min[1] + unsafe_max[1]) / 2.0
    #     p2 = (unsafe_min[2] + unsafe_max[2]) / 2.0
    #     r = p0 - unsafe_min[0]
    #     m.addConstr((x0[0] - p0) * (x0[0] - p0) + (x0[1] - p1) * (x0[1] - p1) + (x0[2] - p2) * (x0[2] - p2) <= r * r)

    f1_x0, df1dx_x0 = barrier_f(m, x0, wb_1_list, wb_length, True)
    # f2_x0, df2dx_x0 = barrier_f(m, x0, wb_2_list, wb_length, True)

    # xt = np.random.randn(DIM)
    xt = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    result_pre = 0

    while True:
        # print("================")
        # f1_xt, df1dx_xt = barrier_f(m, xt, wb_1_list, wb_length, False)
        f2_xt, df2dx_xt = barrier_f(m, xt, wb_2_list, wb_length, False)

        # print("f1_xt:", f1_xt)
        # print("df1dx_xt[0]:", df1dx_xt)
        # print("f2_x0:", f2_x0)

        m.setObjective(f1_x0
                       - f2_xt
                       - df2dx_xt[0] * (x0[0] - xt[0]) - df2dx_xt[1] * (x0[1] - xt[1]) - df2dx_xt[2] * (x0[2] - xt[2])
                       - df2dx_xt[3] * (x0[3] - xt[3]) - df2dx_xt[4] * (x0[4] - xt[4]) - df2dx_xt[5] * (x0[5] - xt[5])
                       - df2dx_xt[6] * (x0[6] - xt[6]) - df2dx_xt[7] * (x0[7] - xt[7]) - df2dx_xt[8] * (x0[8] - xt[8])
                       - df2dx_xt[9] * (x0[9] - xt[9]),
                       GRB.MINIMIZE)
        m.optimize()
        # print("x0.x:", x0.x)
        # print("m.objVal:", m.objVal)
        if abs(m.objVal - result_pre) < margin1 or np.linalg.norm(x0.x - xt) < margin2:
            break
        xt = x0.x
        # print("xt[0]:", xt[0])
        # print("xt[1]:", xt[1])
        result_pre = m.objVal

    # print("minimize unsafe value:", m.objVal)  # 0.36436068000596666
    if m.objVal < 0:
        cex_unsafe.append(np.array(x0.x))
        res_unsafe.append(m.objVal)
    return cex_unsafe, res_unsafe


def verify_thirdCond(wb_1_list, wb_2_list, wb_length, x_len):
    cex_domain = []
    res_domain = []

    data_dir = "./data/"
    pieces = int(4 / x_len)
    datas = np.loadtxt(data_dir + f"-2_2_{piece_len}_20.txt", dtype=np.float64).reshape([-1, 4])
    # datas = np.loadtxt(data_dir + f"03_-0.3_0.3_{piece_len}_20.txt", dtype=np.float64).reshape([-1, 3])

    x0_min = domain_min[0]
    x0_max = domain_max[0]
    x1_min = domain_min[1]
    x1_max = domain_max[1]
    x2_min = domain_min[2]
    x2_max = domain_max[2]
    x3_min = domain_min[3]
    x3_max = domain_max[3]
    x4_min = domain_min[4]
    x4_max = domain_max[4]
    x5_min = domain_min[5]
    x5_max = domain_max[5]
    x6_min = domain_min[6]
    x6_max = domain_max[6]
    x7_min = domain_min[7]
    x7_max = domain_max[7]
    x8_min = domain_min[8]
    x8_max = domain_max[8]
    x9_min = domain_min[9]
    x9_max = domain_max[9]

    m = gp.Model()
    m.setParam('Outputflag', 0)
    # m.setParam('InfUnbdInfo', 1)

    x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x0', lb=-np.inf, ub=np.inf)
    x0[0].lb = x0_min
    x0[0].ub = x0_max
    x0[1].lb = x1_min
    x0[1].ub = x1_max
    x0[2].lb = x2_min
    x0[2].ub = x2_max
    x0[3].lb = x3_min
    x0[3].ub = x3_max
    x0[4].lb = x4_min
    x0[4].ub = x4_max
    x0[5].lb = x5_min
    x0[5].ub = x5_max
    x0[6].lb = x6_min
    x0[6].ub = x6_max
    x0[7].lb = x7_min
    x0[7].ub = x7_max
    x0[8].lb = x8_min
    x0[8].ub = x8_max
    x0[9].lb = x9_min
    x0[9].ub = x9_max


    x1 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x1', lb=-np.inf, ub=np.inf)  # x1 = F(kx+b; x0)
    x1[0].lb = x0_min
    x1[0].ub = x0_max
    x1[1].lb = x1_min
    x1[1].ub = x1_max
    x1[2].lb = x2_min
    x1[2].ub = x2_max
    x1[3].lb = x3_min
    x1[3].ub = x3_max
    x1[4].lb = x4_min
    x1[4].ub = x4_max
    x1[5].lb = x5_min
    x1[5].ub = x5_max
    x1[6].lb = x6_min
    x1[6].ub = x6_max
    x1[7].lb = x7_min
    x1[7].ub = x7_max
    x1[8].lb = x8_min
    x1[8].ub = x8_max
    x1[9].lb = x9_min
    x1[9].ub = x9_max

    m.addConstr(x1[0] == x0[0] + time_step * (x0[3]))
    m.addConstr(x1[1] == x0[1] + time_step * (x0[4]))
    m.addConstr(x1[2] == x0[2] + time_step * (x0[5]))
    m.addConstr(x1[3] == x0[3] + time_step * (-7253.4927 * x0[0] + 1936.3639 * x0[9] - 1338.7624 * x0[3] + 1333.3333 * x0[7]))
    m.addConstr(x1[4] == x0[4] + time_step * (-1936.3639 * x0[8] - 7253.4927 * x0[1] - 1338.7624 * x0[4] - 1333.3333 * x0[6]))
    m.addConstr(x1[5] == x0[5] + time_step * (-769.2308 * x0[2] - 770.2301 * x0[5]))
    m.addConstr(x1[6] == x0[6] + time_step * (x0[8]))
    m.addConstr(x1[7] == x0[7] + time_step * (x0[9]))
    m.addConstr(x1[8] == x0[8] + time_step * (9.81 * x0[1]))
    m.addConstr(x1[9] == x0[9] + time_step * (-9.81 * x0[0]))

    f1_x0, df1dx_x0 = barrier_f(m, x0, wb_1_list, wb_length, True)
    # f2_x0, df2dx_x0 = barrier_f(m, x0, wb_2_list, wb_length, True)
    # f1_x1, df1dx_x1 = barrier_f(m, x1, wb_1_list, wb_length, True)
    f2_x1, df2dx_x1 = barrier_f(m, x1, wb_2_list, wb_length, True)
    # f1_x2, df1dx_x2 = barrier_f(m, x2, wb_1_list, wb_length, True)
    # f2_x2, df2dx_x2 = barrier_f(m, x2, wb_2_list, wb_length, True)

    # xt = np.random.randn(DIM)
    xt = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
    result_pre = -10000

    while True:
        # _1 means use piecewise function  # x0_next_1 == x1
        # x0_next_1 = [x0[0] + time_step * (- x0[0] + (w0 * x0[0] + w1 * x0[1] + var_b)),
        #              x0[1] + time_step * (- x0[1])]

        # _2 means use origin function
        xt_next_2 = [xt[0] + time_step * (xt[3]),
                     xt[1] + time_step * (xt[4]),
                     xt[2] + time_step * (xt[5]),
                     xt[3] + time_step * (-7253.4927 * xt[0] + 1936.3639 * xt[9] - 1338.7624 * xt[3] + 1333.3333 * xt[7]),
                     xt[4] + time_step * (-1936.3639 * xt[8] - 7253.4927 * xt[1] - 1338.7624 * xt[4] - 1333.3333 * xt[6]),
                     xt[5] + time_step * (-769.2308 * xt[2] - 770.2301 * xt[5]),
                     xt[6] + time_step * (xt[8]),
                     xt[7] + time_step * (xt[9]),
                     xt[8] + time_step * (9.81 * xt[1]),
                     xt[9] + time_step * (-9.81 * xt[0])]
        # print("xt_next_2:", xt_next_2)

        # f1_xt, df1dx_xt = barrier_f(m, xt, wb_1_list, wb_length, False)
        f2_xt, df2dx_xt = barrier_f(m, xt, wb_2_list, wb_length, False)
        # print("f2_xt:", f2_xt)
        # print("df2dx_xt:", df2dx_xt)

        f1_xt_next, df1dx_xt_next = barrier_f(m, xt_next_2, wb_1_list, wb_length, False)
        # print("f1_xt_next:", f1_xt_next)
        # print("df1dx_xt_next:", df1dx_xt_next)
        # f2_xt_next, df2dx_xt_next = barrier_f(m, xt_next_2, wb_2_list, wb_length, False)

        m.setObjective(f1_xt_next
                       + df1dx_xt_next[0] * (x1[0] - xt_next_2[0]) + df1dx_xt_next[1] * (x1[1] - xt_next_2[1])
                       + df1dx_xt_next[2] * (x1[2] - xt_next_2[2]) + df1dx_xt_next[3] * (x1[3] - xt_next_2[3])
                       + df1dx_xt_next[4] * (x1[4] - xt_next_2[4]) + df1dx_xt_next[5] * (x1[5] - xt_next_2[5])
                       + df1dx_xt_next[6] * (x1[6] - xt_next_2[6]) + df1dx_xt_next[7] * (x1[7] - xt_next_2[7])
                       + df1dx_xt_next[8] * (x1[8] - xt_next_2[8]) + df1dx_xt_next[9] * (x1[9] - xt_next_2[9])
                       - f2_x1
                       - (1 - alpha) * f1_x0
                       + (1 - alpha) * (f2_xt + df2dx_xt[0] * (x0[0] - xt[0]) + df2dx_xt[1] * (x0[1] - xt[1])
                                        + df2dx_xt[2] * (x0[2] - xt[2]) + df2dx_xt[3] * (x0[3] - xt[3])
                                        + df2dx_xt[4] * (x0[4] - xt[4]) + df2dx_xt[5] * (x0[5] - xt[5])
                                        + df2dx_xt[6] * (x0[6] - xt[6]) + df2dx_xt[7] * (x0[7] - xt[7])
                                        + df2dx_xt[8] * (x0[8] - xt[8]) + df2dx_xt[9] * (x0[9] - xt[9])),
                       GRB.MAXIMIZE)
        m.optimize()

        if m.status == 2:
            if m.objVal < result_pre:
                X = xt
                Val = result_pre
                break
            if abs(m.objVal - result_pre) < margin1 or np.linalg.norm(x0.x - xt) < margin2:
                X = x0.x
                Val = m.objVal
                break
            xt = x0.x
            # print("x0.x:", x0.x)
            # print("m.objVal:", m.objVal)
            result_pre = m.objVal
        else:
            break
    # print(f"{cex_idx} is solvered!")
    if m.status == 2:
        # print("maximize domain value:", m.objVal)
        if Val > 0:
            cex_domain.append(np.array(X))
            res_domain.append(Val)
        # else:
            # if x0 is a boundary point, x1 and x2 may leave domain set (violate with x1.lb/x1.ub/x2.lb/x2.ub)
            # if no one point satisfy these bound condition, solver cant get objVal, we can still say this block is safe
            # print("not converge")

    return cex_domain, res_domain


def verif(model, x_len):
    wb_1_list = []
    wb_2_list = []
    params = model.state_dict()
    i = 0
    for k in params:
        # 'hidden_layers1.0.weight' 'hidden_layers1.0.bias' 'hidden_layers2.0.weight' 'hidden_layers2.0.bias'
        if 'layers' in k:
            continue
        if i % 4 < 2:
            wb_1_list.append(params[k].detach().numpy())
        else:
            wb_2_list.append(params[k].detach().numpy())
        i += 1

    wb_length = int(len(wb_1_list) / 2)

    cex_init, res_init = verify_initCond(wb_1_list, wb_2_list, wb_length)
    cex_unsafe, res_unsafe = verify_unsafeCond(wb_1_list, wb_2_list, wb_length)
    cex_domain, res_domain = verify_thirdCond(wb_1_list, wb_2_list, wb_length, x_len)
    print("cex_init: ", cex_init)
    print("res_init: ", res_init)
    print("cex_unsafe: ", cex_unsafe)
    print("res_unsafe: ", res_unsafe)
    print("cex_domain: ", cex_domain)
    print("res_domain: ", res_domain)
    return cex_init, cex_unsafe, cex_domain


if __name__ == "__main__":
    model = cdinn2.gen_nn()
    model.load_state_dict(torch.load('../model/cdinn_zhrc1_1_20_0_0.1_1_train.pt'), strict=True)
    verif(model, piece_len)
